# encoding: utf-8
import os
import shutil
import re
import time
from functools import wraps

import torch
from torch.mps import empty_cache


def get_offset(text, target, default_start_index=0, default_end_index=None):
    """
    获取字符串特定值位置集合
    :param text:
    :param target:
    :param default_start_index:
    :param default_end_index:
    :return:
    """
    offsets = []
    # 初始化起始索引
    start_index = default_start_index
    end_index = default_end_index if default_end_index is not None else len(text)
    while start_index != -1:
        start_index = text.find(target, start_index, end_index)
        if start_index != -1:
            offsets.append([start_index, start_index + len(target)])
            start_index += len(target)  # 跳过当前匹配的 "sf"
    return offsets


def rm_file(target_directory):

    # 检查目录是否存在
    if os.path.exists(target_directory):
        # 删除目录及其内容
        shutil.rmtree(target_directory)
        # print(f'目录 "{target_directory}" 已被删除。')
    else:
        print(f'目录 "{target_directory}" 不存在。')


def run_time(func):
    @wraps(func)
    def wrapper(*args, **kwargs):
        start_time = time.time()
        result = func(*args, **kwargs)
        end_time = time.time()
        print(f"程序 {func.__name__} 运行时间: {end_time - start_time} 秒")
        return result
    return wrapper


def torch_gc():
    try:
        if torch.cuda.is_available():
            # with torch.cuda.device(DEVICE):
            torch.cuda.empty_cache()
            torch.cuda.ipc_collect()
        elif torch.backends.mps.is_available():
            try:
                empty_cache()
            except Exception as e:
                msg = ("如果您使用的是 macOS 建议将 pytorch 版本升级至 2.0.0 或更高版本，"
                       "以支持及时清理 torch 产生的内存占用。")
    except Exception:
        ...


def parse_value_str(value: str):
    try:
        value = int(value)
    except:
        try:
            value = float(value)
        except:
            pass
    return value


def parse_llm_json(llm_json: str):
    llm_dict = {}
    json_str_split = re.findall(r'"(.*?)": (.*?)[,\n]', llm_json)
    for key, value in json_str_split:
        key: str
        value: str
        if value.startswith('"'):
            value = value[1:]
        if value.endswith('"'):
            value = value[:-1]
        llm_dict[key] = parse_value_str(value)
    return llm_dict


def split_list(lst, n):
    return [lst[i:i + n] for i in range(0, len(lst), n)]


if __name__ == '__main__':
    find_text = 'abacdef, defdd'
    aaa = get_offset(find_text, 'd', 3, len(find_text) - 1)
    print(aaa)